# -*- coding: utf-8 -*-
# created on 2016/3/8

from sympy import Add, Mul, Eq, sympify, real_root, Poly, latex, factor_terms, Pow, sqrt, solveset, S, log, Rel
from mathsolver.functions.sympy_utils import get_all_child, get_pow, simple_invert


def _sqrt(d):
    # remove squares from square root since both will be represented
    # in the results; a similar thing is happening in roots() but
    # must be duplicated here because not all quadratics are binomials
    co = []
    other = []
    for di in Mul.make_args(d):
        if di.is_Pow and di.exp.is_Integer and di.exp % 2 == 0:
            co.append(Pow(di.base, di.exp // 2))
        else:
            other.append(di)
    if co:
        d = Mul(*other)
        co = Mul(*co)
        return co * sqrt(d)
    return sqrt(d)


ZHENGLI = "整理方程："
HUAJIAN = "等式两边同除以%s："
INVERT = "对%s函数求逆："
LOG = "取log："
FORMULA = "使用求根公式："
FACTOR = "因式分解："
UNION = "解得："
POW_INVERT1 = "求解偶数幂："
POW_INVERT2 = "求解奇数幂："
SUB = "对(%d)，"
WUJIE = "无解"
INFINITE = "没有有限解"
zero = sympify("0")


class StepSolver(object):
    def __init__(self, f1, f2):
        self.eqs = [((f1, f2), 0)]
        self.f1 = f1
        self.f2 = f2
        self.steps = []
        self.symbol = None
        self.index = 1
        self.check()
        self.answers = []
        self.prefix = ""
        self.ineq = []
        self.solveset = None

    def check(self):
        free_symbols = Eq(self.f1, self.f2).free_symbols
        if len(free_symbols) == 1:
            self.symbol = free_symbols.pop()
        else:
            raise ValueError("Use one variable only!")
        try:
            self.solveset = solveset(Eq(self.f1, self.f2), self.symbol)
        except Exception:
            pass

    def steps_append(self, text, eqs):
        self.steps.append((self.prefix + text, eqs))
        self.prefix = ""

    def eqs_append(self, eqs, text):
        new_eqs = []
        if len(eqs) == 1:
            new_eqs.append((eqs[0], 0))
        else:
            for eq in eqs:
                new_eqs.append((eq, self.index))
                self.index += 1

        self.eqs.extend(new_eqs)
        self.steps_append(text, new_eqs)

    def _solve(self, f1, f2):
        if f2.is_number and f1.is_Function:  # log(..) = n || exp(..) = n
            f1, f2 = simple_invert(f1, f2)
            return self.eqs_append([(f1, f2)], INVERT % f1.inverse)

        if f2.is_zero and f1.is_Mul:  # .. = 0
            return self.simple_mul(f1)

        factor = f1.factor()
        if f2.is_zero and factor != f1 and factor.is_Mul:  # () * () = 0
            self.steps_append(FACTOR, [((factor, zero), 0)])
            return self.simple_mul(factor)
        elif f2.is_zero and factor != f1 and factor.is_Pow:
            self.steps_append(FACTOR, [((factor, zero), 0)])
            return self.solve_pow(factor, zero)

        if f1.is_Pow and f1.args[0].has(self.symbol):  # x^k = n
            return self.solve_pow(f1, f2)

        if f2.is_zero and f1.is_Add:  # x^2 + 2x + 1 = 0
            pows = [get_pow(arg, self.symbol) for arg in f1.args]
            if set(pows) == {0, 1, 2}:
                return self.solve_quadratic(f1)

        if f2.is_number and f1.is_Mul:  # 9*x^2 = 1
            a, h = f1.as_independent(self.symbol)
            return self.eqs_append([(h, f2 / a)], HUAJIAN % ("%.5f" % a).rstrip("0"))

        f = Add(f1, -f2).expand()
        if len(get_all_child(f, lambda x: x.is_Pow and x.args[1] == -1 and x.args[0].has(self.symbol))) > 1:
            f = f.together()
            fenmu = get_all_child(f, lambda x: x.is_Pow and x.args[1] == -1 and x.args[0].has(self.symbol))
            fenmu = Mul._from_args(fenmu)
            for ans in solveset(1 / fenmu, self.symbol):
                self.ineq.append(Rel(self.symbol, ans, "!="))
            return self.eqs_append([((f * (1 / fenmu)).simplify(), zero)], "去分母得")

        if f.is_Add:
            a, h = f.as_independent(self.symbol)
            pows = [get_pow(arg, self.symbol) for arg in f.args]
            new_factor = f.factor()
            if h.is_Function or h.is_Pow or h.is_Mul:
                return self.eqs_append([(h, -a)], ZHENGLI)
            elif set(pows) == {0, 1, 2} or new_factor.is_Mul or new_factor.is_Pow:
                return self.eqs_append([(f, zero)], ZHENGLI)
            else:
                return self.eqs_append([(h, -a)], ZHENGLI)
        else:
            raise ValueError("Unexpected Situation!")

    def simple_mul(self, f1):  # f1 = 0
        eqs = []
        for arg in f1.args:
            if arg.has(self.symbol) and not (arg.is_Pow and arg.args[1] < 0):
                eqs.append((arg, zero))

        self.eqs_append(eqs, UNION)

    def solve_pow(self, f1, f2):  # f1 is Pow, f2 is number
        base, _pow = f1.args
        if base.has(self.symbol) and _pow.has(self.symbol):
            raise ValueError("Too Complex!")
        elif base.has(self.symbol):  # x^k = n
            pow_num = f1.args[1]
            res = real_root(f2, pow_num)
            eqs = [(base, res)]
            if pow_num.is_even:
                if not res.is_zero:
                    eqs.append((base, -res))
                self.eqs_append(eqs, POW_INVERT1)
            else:
                self.eqs_append(eqs, POW_INVERT2)

            if f2.has(self.symbol) and (pow_num.is_even or (1 / pow_num).is_even):
                self.ineq.append(Rel(f2, S.Zero, ">"))

        elif base.is_Number:  # k^x = n
            res = log(f2, base)
            self.eqs_append((_pow, res), LOG)
        else:
            raise ValueError("Unexpected Situation!")

    def solve_quadratic(self, f1):  # a*x^2 + b*x + c = 0
        poly = Poly(f1.as_expr())
        a, b, c = poly.all_coeffs()
        if b > 0:
            formula = "%s = {-%s \\pm \\sqrt{%s^2-4*%s*%s} \\over 2*%s}" % (self.symbol, b, b, a, c, a)
        elif b < 0:
            formula = "%s = {%s \\pm \\sqrt{%s^2-4*%s*%s} \\over 2*%s}" % (self.symbol, -b, b, a, c, a)
        self.steps_append(FORMULA, [(formula, 0)])

        d = b ** 2 - 4 * a * c
        da = 2 * a
        db = -b / da

        dd = factor_terms(_sqrt(d) / da)
        r0 = db - dd
        r1 = db + dd
        if a.is_negative:
            r0, r1 = r1, r0

        self.eqs_append([(self.symbol, r0), (self.symbol, r1)], "")

    def check_answer(self):
        solves = set(self.solveset.args)
        answers = set(self.answers)

        if not self.solveset.is_EmptySet:
            if solves.intersection(answers) == solves:
                return False
            elif len(answers.difference(solves)) > 0:
                raise ValueError("Get an unexpected answer!")

        else:
            if len(answers) > 0:
                if all([ineq.subs({self.symbol: ans}) is False for ans in answers for ineq in self.ineq]):
                    return False

        return True

    def solve(self):
        num = 0
        try:
            while self.check_answer() and len(self.eqs) > 0 and num < 20:
                (f1, f2), index = self.eqs.pop()
                if f1 == self.symbol and not f2.has(self.symbol):
                    if all([ineq.subs({self.symbol: f2}) for ineq in self.ineq]):
                        self.answers.append(f2)
                else:
                    if index:
                        self.prefix = SUB % index
                    self._solve(f1, f2)

                    new_f1, new_f2 = self.eqs[-1][0]
                    if f1 == new_f1 and f2 == new_f2:
                        self.eqs.pop()
                num += 1
        except Exception:
            pass

        return self

    def printing(self):
        print_steps = []  # [["", "$$%s$$" % latex(Eq(self.f1, self.f2))]]
        if not self.solveset.is_FiniteSet:
            print_steps.append([INFINITE, ""])
            return print_steps

        last_eqs = ""
        for text, eqs in self.steps:
            if eqs == last_eqs:
                continue

            if len(eqs) == 1:
                eq, index = eqs[0]
                if isinstance(eq, str):
                    eq_str = "$$%s$$" % eq
                else:
                    eq_str = "$$%s$$" % latex(Eq(*eq))

            elif len(eqs) > 1:
                eq_list = []
                for eq, index in eqs:
                    eq = latex(Eq(*eq))
                    if index:
                        eq += "\\quad(%d)" % index
                    eq_list.append(eq)

                eq_str = "\\\\".join(eq_list)
                eq_str = "$$\\left\\{\\begin{array}{c}%s\\end{array}\\right.$$" % eq_str

            print_steps.append(["", text + eq_str])
            last_eqs = eqs

        answer = []
        for ans in self.answers:
            _str = "%s = %s" % (self.symbol, latex(ans))
            answer.append(_str)

        if len(answer) > 1:
            answer = "\\\\".join(answer)
            answer = "$$\\left\\{\\begin{array}{c}%s\\end{array}\\right.$$" % answer
        elif len(answer) == 1:
            answer = "$$%s$$" % answer[0]

        if not self.solveset.is_EmptySet:
            print_steps.append(["", "∴" + answer])
        else:
            print_steps.append(["$$因为x \\ne %s$$" % ",".join([latex(v) for v in self.ineq]), "无解"])
        return print_steps


if __name__ == '__main__':
    pass
